#pragma once
#include <iostream>
#include <memory>

#include <cuda_runtime.h>

#include ""

namespace bl
{

struct DeviceTags
{
    struct CPU;
    struct CUDA;
};

template <typename T>
T* allocateCUDA(int width, int height, size_t& pitch)
{
	T* device_ptr;
	cudaMallocPitch((void**)&device_ptr, &pitch, width * sizeof(T), height);

	return device_ptr;
}

template <typename T>
T* allocateCPU(int width, int height, size_t& pitch)
{
	T* host_ptr = new T[width * height];
    pitch       = width * sizeof(T);

	return host_ptr;
}

template <typename T>
void freeCUDA(T* device_ptr)
{
    if (device_ptr != nullptr)
        cudaFree(device_ptr);
}

template <typename T>
void freeCPU(T* host_ptr)
{
    if (host_ptr != nullptr)
        delete[] host_ptr;
}

template <typename T, typename Device>
class Allocater;

template <typename T>
class Allocater<T, DeviceTags::CUDA>
{
private:
    std::shared_ptr<T>      __ptr{nullptr};
    size_t                  __linesize{0};

public:
    Allocater(int width, int height)
    {
        __ptr = std::shared_ptr<T>(allocateCUDA<T>(width, height, __linesize), [](T* ptr) { freeCUDA(ptr); });
    }
};

template <typename T>
class Allocater<T, DeviceTags::CPU>
{
private:
    std::shared_ptr<T>      __ptr{nullptr};
    size_t                  __linesize{0};

public:
    Allocater(int width, int height)
    {
        __ptr = std::shared_ptr<T>(allocateCPU<T>(width, height, __linesize), [](T* ptr) { freeCPU(ptr); });
    }
};

}